import os
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
from src.llamafactory.chat.chat_model import ChatModel
from tqdm import tqdm
import argparse
import json

output = '/root/autodl-tmp/sft_1209_rank64_train0316_1e5_4_5epoch_seed8888'
parser = argparse.ArgumentParser(description="Your script description")
parser.add_argument("--model_name_or_path", default="/root/autodl-tmp/Qwen2-7B-Instruct", type=str,
                    help="Model name or path")
parser.add_argument("--adapter_name_or_path", default=output, type=str, help="Checkpoint directory")
parser.add_argument("--template", default="qwen", type=str, help="Template type")
parser.add_argument("--finetuning_type", default="lora", type=str, help="Fine-tuning type")
# parser.add_argument("--quantization_bit", default=4, type=str, help="quantization_bit")

args = parser.parse_args()
args_dict = vars(args)
chat_model = ChatModel(args_dict)
history = []



import json
import re

with open('./合成数据决赛赛题.jsonl','r') as f:
    data = f.readlines()
data[0]

# datas=[]
# for t in test:
#     t=json.loads(t)
#     datas.append(t)
fout = open(f"{output}/submission_testb_all_5epoch0316.json", "w", encoding="utf-8")
results1 = []
new_results=[]

num=0
for d in tqdm(data):
    # text = df['正文前 256符号'][i]
    d = json.loads(d)
    apis = d['apis']
    user_message = d['user_messages']

    prompt = f"""你是一名专业的智能体agent执行者，作为agent执行任务的过程中，通常会根据用户的多轮对话内容，总结问题，使用的相应工具来完成任务，给出使用相应工具api的名称和参数。 
    ##要求： 
    1.用户的多轮对话中可能需要调用多个api工具以及一个工具api调用多次，需要给出工具名 
    2.严格按照提供的工具agent的参数提供内容, api的required参数都需要根据多轮对话内容抽出匹配的结果 
    3.如果required参数中有空则整个api参数arguments为空字典 
    4.如果api的非required参数为空，则arguments中不出现该参数 
    5.应该根据当前多轮对话内容，需要判断出在第几轮内容可以满足工具调用的参数需求，如果在第三轮可以判断出工具调用的api以及所需要的参数，则在第三轮给出api调用，如果在第四轮则放在第四轮 请根据下面的输入，输出正确的结果
    6.在用户的多轮对话中可能会存在一个api调用多次，也可能会存在多个api调用，以及只有一次api调用，需要你在每轮对话的结束给出合理的答案。
    
    请根据下面的输入，输出正确的结果：
    agents:
    {apis}
    用户多轮对话内容:
    {user_message}
    """
   
    messages = [
        {"role": "user", "content":   prompt}]

    response = chat_model.chat(messages, history=[],  top_p=0.8, temperature=0.1,do_sample=False,max_new_tokens=2048,num_beams=3)
    
    res = response[0].response_text.strip()
    
    # try:
    #     res  = eval(res.replace('true','True').replace('false','False'))
    # except:
    #     num+=1
    #     res = []
    print(res)
    results1.append(res)
    d['targets']=res
    new_results.append(d) 
    fout.write(json.dumps(d, ensure_ascii=False) + "\n")
    fout.flush()   
# print(num/len(test))
with open(f'{output}/results_7b_5epoch.json', 'w') as file:
       json.dump(new_results, file, indent=4, ensure_ascii=False)